import time
from typing_extensions import TypedDict
from langgraph.graph import StateGraph
from langgraph.cache.memory import InMemoryCache
from langgraph.types import CachePolicy


class State(TypedDict):
    x: int
    result: int

def expensive_node(state: State) -> dict[str, int]:
    # expensive computation
    time.sleep(2)
    return {"result": state["x"] * 2}


if __name__ == '__main__':
    builder = StateGraph(State)
    builder.add_node("expensive_node", expensive_node, cache_policy=CachePolicy(ttl=1))
    # builder.add_node("expensive_node", expensive_node, cache_policy=CachePolicy(ttl=1))
    builder.set_entry_point("expensive_node")
    builder.set_finish_point("expensive_node")

    graph = builder.compile(cache=InMemoryCache())

    print(graph.invoke({"x": 5}, stream_mode='updates'))
    time.sleep(2)
    # [{'expensive_node': {'result': 10}}]
    print(graph.invoke({"x": 5}, stream_mode='updates'))
    # [{'expensive_node': {'result': 10}, '__metadata__': {'cached': True}}]